import argparse

import torch
import torch.nn as nn
from torch import optim
import torchvision
from torchvision import models,datasets,transforms
import os
from tqdm import tqdm
import h5py
from torch.utils.data import DataLoader
#—————————————————————————————————— 流程，功能   ————————————————————————————————— #
"""
1、准备数据
2、建立提取数据的网络结构
3、建立预测网络结构
4、开始训练
4.1 准备参数
"""
#—————————————————————————————————— 1、准备参数 ————————————————————————————————— #
use_gpu = torch.cuda.is_available()
epochs = 5
batch_size = 32
model_list = ["resnet18","vgg"]
parse = argparse.ArgumentParser()
parse.add_argument(
    '--model',default="vgg", help='vgg, inceptionv3, resnet152')
parse.add_argument('--phase', help='train, val')
opt = parse.parse_args()


#—————————————————————————————————— 2、准备数据 ————————————————————————————————— #
train_dataset = torchvision.datasets.CIFAR10('D:\code\datasets\cifar-10\cifar_code\data', train=True, download=False,
                                        transform=torchvision.transforms.Compose([
                                         torchvision.transforms.ToTensor(), # 将图像数据从PIL类型变换成32位浮点数格式
                                         torchvision.transforms.Normalize((0.1307,), (0.3081,)),]))

validation_dataset = torchvision.datasets.CIFAR10('D:\code\datasets\cifar-10\cifar_code\data', train=False, download=False,
                                        transform=torchvision.transforms.Compose([
                                         torchvision.transforms.ToTensor(), # 将图像数据从PIL类型变换成32位浮点数格式
                                         torchvision.transforms.Normalize((0.1307,), (0.3081,)),]))

# Create train and validation batch for training
training_loader = torch.utils.data.DataLoader(
    dataset=train_dataset, batch_size=batch_size, shuffle=False,num_workers=1)
validation_loader = torch.utils.data.DataLoader(
    dataset=validation_dataset, batch_size=batch_size,num_workers=1)
dataloader = {
    'train':
    training_loader,
    'val':
    validation_loader
}

#—————————————————————————————————— 3、模型训练————————————————————————————————— #

# 3.1 准备提取图像特征的网络
class feature_net(nn.Module):
    def __init__(self, model):
        super(feature_net, self).__init__()
        if model == 'vgg':
            vgg = models.vgg19(pretrained=True)

            self.feature = nn.Sequential(*list(vgg.children())[:-1])
            # 这一步输出结果为（B,512,7,7）,512*7*7=25088
            self.feature.add_module('global average', nn.AvgPool2d(7))
            # 经过最大池化层，结果输出为（512,1,1）
        elif model == 'resnet18':
            resnet = models.resnet18(pretrained=True)
            print(resnet)
            self.feature = nn.Sequential(*list(resnet.children())[:-1])
            print(self.feature)

    def forward(self, x):
        """
        model includes vgg19, inceptionv3, resnet152
        """
        x = self.feature(x)
        x = x.view(x.size(0), -1)
        return x

class classifier(nn.Module):
    # 全连接分类层
    def __init__(self, dim, n_classes):
        super(classifier, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(dim, 1000),
            nn.ReLU(True),
            nn.Dropout(0.5),
            nn.Linear(1000, n_classes)
        )

    def forward(self, x):
        x = self.fc(x)
        return x

# 3.2 提取特征
def CreateFeature(model,phase, outputPath='.'):
    featurenet = feature_net(model).cuda()

    if use_gpu:
        featurenet.cuda()
    feature_map = torch.FloatTensor()
    label_map = torch.LongTensor()

    for data in tqdm(dataloader["train"]):
        img, label = data
        if use_gpu:
            img = img.cuda()
        out = featurenet(img)
        feature_map = torch.cat((feature_map, out.cpu().data), 0)
        # dim=0，表示按列合并，在行的维度增加数据
        label_map = torch.cat((label_map, label), 0)
        # 一直在后面增加标签，
    feature_map = feature_map.numpy()
    label_map = label_map.numpy()
    file_name = '_feature_{}.hd5f'.format(model)
    #
    h5_path = os.path.join(outputPath, phase) + file_name
    with h5py.File(h5_path, 'w') as h:
        h.create_dataset('data', data=feature_map)
        h.create_dataset('label', data=label_map)


if __name__ == '__main__':
    for i in model_list:
        CreateFeature(i,opt.phase)



